import pandas as pd
import numpy as np
import math
from collections import Counter

data_table = pd.DataFrame([
    {'色泽': '青绿', '根蒂': '蜷缩', '敲声': '浊响', '纹理': '清晰', '脐部': '凹陷', '触感': '硬滑', '密度': 0.697, '含糖量': 0.460, '好瓜': '是'},
    {'色泽': '乌黑', '根蒂': '蜷缩', '敲声': '沉闷', '纹理': '清晰', '脐部': '凹陷', '触感': '硬滑', '密度': 0.774, '含糖量': 0.376, '好瓜': '是'},
    {'色泽': '乌黑', '根蒂': '蜷缩', '敲声': '浊响', '纹理': '清晰', '脐部': '凹陷', '触感': '硬滑', '密度': 0.634, '含糖量': 0.264, '好瓜': '是'},
    {'色泽': '青绿', '根蒂': '蜷缩', '敲声': '沉闷', '纹理': '清晰', '脐部': '凹陷', '触感': '硬滑', '密度': 0.608, '含糖量': 0.318, '好瓜': '是'},
    {'色泽': '浅白', '根蒂': '蜷缩', '敲声': '浊响', '纹理': '清晰', '脐部': '凹陷', '触感': '硬滑', '密度': 0.556, '含糖量': 0.215, '好瓜': '是'},
    {'色泽': '青绿', '根蒂': '稍蜷', '敲声': '浊响', '纹理': '清晰', '脐部': '稍凹', '触感': '软粘', '密度': 0.403, '含糖量': 0.237, '好瓜': '是'},
    {'色泽': '乌黑', '根蒂': '稍蜷', '敲声': '浊响', '纹理': '稍糊', '脐部': '稍凹', '触感': '软粘', '密度': 0.481, '含糖量': 0.149, '好瓜': '是'},
    {'色泽': '乌黑', '根蒂': '稍蜷', '敲声': '浊响', '纹理': '清晰', '脐部': '稍凹', '触感': '硬滑', '密度': 0.437, '含糖量': 0.211, '好瓜': '是'},
    {'色泽': '乌黑', '根蒂': '稍蜷', '敲声': '沉闷', '纹理': '稍糊', '脐部': '稍凹', '触感': '硬滑', '密度': 0.666, '含糖量': 0.091, '好瓜': '否'},
    {'色泽': '青绿', '根蒂': '硬挺', '敲声': '清脆', '纹理': '清晰', '脐部': '平坦', '触感': '软粘', '密度': 0.243, '含糖量': 0.267, '好瓜': '否'},
    {'色泽': '浅白', '根蒂': '硬挺', '敲声': '清脆', '纹理': '模糊', '脐部': '平坦', '触感': '硬滑', '密度': 0.245, '含糖量': 0.057, '好瓜': '否'},
    {'色泽': '浅白', '根蒂': '蜷缩', '敲声': '浊响', '纹理': '模糊', '脐部': '平坦', '触感': '软粘', '密度': 0.343, '含糖量': 0.099, '好瓜': '否'},
    {'色泽': '青绿', '根蒂': '稍蜷', '敲声': '浊响', '纹理': '稍糊', '脐部': '凹陷', '触感': '硬滑', '密度': 0.639, '含糖量': 0.161, '好瓜': '否'},
    {'色泽': '浅白', '根蒂': '稍蜷', '敲声': '沉闷', '纹理': '稍糊', '脐部': '凹陷', '触感': '硬滑', '密度': 0.657, '含糖量': 0.198, '好瓜': '否'},
    {'色泽': '乌黑', '根蒂': '稍蜷', '敲声': '浊响', '纹理': '清晰', '脐部': '稍凹', '触感': '软粘', '密度': 0.360, '含糖量': 0.370, '好瓜': '否'},
    {'色泽': '浅白', '根蒂': '蜷缩', '敲声': '浊响', '纹理': '模糊', '脐部': '平坦', '触感': '硬滑', '密度': 0.593, '含糖量': 0.042, '好瓜': '否'},
    {'色泽': '青绿', '根蒂': '蜷缩', '敲声': '沉闷', '纹理': '稍糊', '脐部': '稍凹', '触感': '硬滑', '密度': 0.719, '含糖量': 0.103, '好瓜': '否'}
], columns=['色泽','根蒂','敲声','纹理','脐部','触感','密度','含糖量','好瓜'])
att = {'色泽': 0, '根蒂': 0, '敲声':0, '纹理': 0, '触感': 0}

#计算信息熵
def entropy(data):
    counter = Counter(data['好瓜'])
    s = sum(counter.values())
    return sum([-p/s * math.log(p/s, 2) for p in counter.values()])

#计算信息增益
def info_gain(data, attr, ent):
    counter = Counter(data[attr])
    gain = 0
    for i in counter:
        gain += entropy(data.loc[data[attr]==i]) * counter[i] / len(data)
        # print(gain)
    return ent - gain 

def tree(root):
    return [root, []]

def node(m_root, data, att):
    counter = Counter(data[m_root[0]])
    fu = Counter(data['好瓜'])
    for i in counter:
        data_d = data.loc[data[m_root[0]] == i]
        counter1 = Counter(data_d['好瓜'])
        #空集
        if len(data_d) == 0:
            t = max(fu)
            if t == '是':
                m_root[-1].append(m_root[0]+i+'好瓜1')
            else:
                m_root[-1].append(m_root[0]+i+'坏瓜1')
            continue
        #无法划分
        if len(att) == 0:
            t = max(counter1)
            if t == '是':
                m_root[-1].append(m_root[0]+i+'好瓜2')
            else:
                m_root[-1].append(m_root[0]+i+'坏瓜2')
            continue
        #全部是一个类
        if counter1['是'] == 0:
            m_root[-1].append(m_root[0]+i+'坏瓜3')
            continue
        elif counter1['否'] == 0:
            m_root[-1].append(m_root[0]+i+'好瓜3')
            continue
        ent = entropy(data_d)
        gain = -1
        root = ''
        for j in att.keys():
            if gain < info_gain(data_d, j, ent):
                root = j
                gain = info_gain(data_d, j, ent)
        m_root[-1].append([])
        m_root[-1][-1].append(root)
        m_root[-1][-1].append([])
        att1 = att.copy()
        att1.pop(root)
        node(m_root[-1][-1], data_d, att1)
    
def TreeGenerator(att, data_table):
    ent = entropy(data_table)
    gain = -1
    root = ''
    for i in att.keys():
        if gain < info_gain(data_table, i, ent):
            root = i
            gain = info_gain(data_table, i, ent)
    Atree = tree(root)
    att.pop(root)
    node(Atree, data_table, att)
    return Atree

a = TreeGenerator(att, data_table)
print(a)
